package com.xs.langchain4j_springboot.splitter;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.document.splitter.DocumentByLineSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.internal.Utils;
import dev.langchain4j.internal.ValidationUtils;
import dev.langchain4j.model.Tokenizer;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

public class CustomizedDocumentByLineSplitter extends DocumentByLineSplitter {
    public CustomizedDocumentByLineSplitter(int maxSegmentSizeInChars, int maxOverlapSizeInChars) {
        super(maxSegmentSizeInChars, maxOverlapSizeInChars, (Tokenizer)null, (DocumentSplitter)null);
    }

    public List<TextSegment> split(Document document) {
        ValidationUtils.ensureNotNull(document, "document");
        List<TextSegment> segments = new ArrayList();
        CustomizedSegmentBuilder segmentBuilder = new CustomizedSegmentBuilder(this.maxSegmentSize, this::estimateSize, this.joinDelimiter());
        AtomicInteger index = new AtomicInteger(0);
        String[] parts = this.split(document.text());
        String overlap = null;

        for(String part : parts) {
            int partSize = segmentBuilder.sizeOf(part);
//            if (segmentBuilder.isNotEmpty()) {
//                String segmentText = segmentBuilder.toString();
//                if (!segmentText.equals(overlap)) {
//                    segments.add(createSegment(segmentText, document, index.getAndIncrement()));
//                    overlap = this.overlapFrom(segmentText);
//                    segmentBuilder.reset();
//                    segmentBuilder.append(overlap);
//                    if (segmentBuilder.hasSpaceFor(partSize)) {
//                        segmentBuilder.append(part);
//                        continue;
//                    }
//                }
//            }
            segmentBuilder.append(part);

            for(TextSegment segment : this.subSplitter.split(Document.from(segmentBuilder.toString()))) {
                segments.add(createSegment(segment.text(), document, index.getAndIncrement()));
            }

            TextSegment lastSegment = (TextSegment)segments.get(segments.size() - 1);
            segmentBuilder.reset();
        }

        if (segmentBuilder.isNotEmpty() && !segmentBuilder.toString().equals(overlap)) {
            segments.add(createSegment(segmentBuilder.toString(), document, index.getAndIncrement()));
        }

        return segments;
    }

    int estimateSize(String text) {
        return this.tokenizer != null ? this.tokenizer.estimateTokenCountInText(text) : text.length();
    }

    static TextSegment createSegment(String text, Document document, int index) {
        Metadata metadata = document.metadata().copy().put("index", String.valueOf(index));
        return TextSegment.from(text, metadata);
    }


}
